import triton
import triton.language as tl
import torch
import torch_npu

from torch._inductor.runtime.triton_helpers import math as tl_math

from .util import triton_helper_compile, compiled_kernel_list

# triton_poi_fused_cat_1
@triton.jit
def fused_kernel8(in_ptr0, in_ptr1, in_ptr2, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 5120
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = (xindex // 64) % 10
    x0 = xindex % 64
    x2 = (xindex // 640)
    x4 = xindex
    tmp0 = x1
    tmp1 = tl.full([1], 0, tl.int32)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 9, tl.int32)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.load(in_ptr0 + (x0 + (64*x2) + (512*x1)), tmp4 & xmask, other=0.0)
    tmp6 = tmp0 >= tmp3
    tmp7 = tl.full([1], 10, tl.int32)
    tmp8 = tmp0 < tmp7
    tmp9 = tl.load(in_ptr1 + (x0 + (64*x2)), tmp6 & xmask, eviction_policy='evict_last', other=0.0)
    tmp10 = tl.load(in_ptr2 + (x4 % 32), tmp6 & xmask, eviction_policy='evict_last', other=0.0)
    tmp11 = tl_math.cos(tmp10)
    tmp12 = 1.0
    tmp13 = tmp11 * tmp12
    tmp14 = tmp9 * tmp13
    tmp15 = x0
    tmp16 = tmp15 >= tmp1
    tmp17 = tl.full([1], 32, tl.int32)
    tmp18 = tmp15 < tmp17
    tmp19 = tmp18 & tmp6
    tmp20 = tl.load(in_ptr1 + (32 + (64*x2) + x0), tmp19 & xmask, eviction_policy='evict_last', other=0.0)
    tmp21 = -tmp20
    tmp22 = tl.full(tmp21.shape, 0.0, tmp21.dtype)
    tmp23 = tl.where(tmp19, tmp21, tmp22)
    tmp24 = tmp15 >= tmp17
    tmp25 = tl.full([1], 64, tl.int32)
    tmp26 = tmp15 < tmp25
    tmp27 = tmp24 & tmp6
    tmp28 = tl.load(in_ptr1 + ((64*x2) + ((-32) + x0)), tmp27 & xmask, eviction_policy='evict_last', other=0.0)
    tmp29 = tl.where(tmp18, tmp23, tmp28)
    tmp30 = tl_math.sin(tmp10)
    tmp31 = tmp30 * tmp12
    tmp32 = tmp29 * tmp31
    tmp33 = tmp14 + tmp32
    tmp34 = tl.full(tmp33.shape, 0.0, tmp33.dtype)
    tmp35 = tl.where(tmp6, tmp33, tmp34)
    tmp36 = tl.where(tmp4, tmp5, tmp35)
    tl.store(out_ptr0 + (x4), tmp36, xmask)

def test_compile_for_fused_kernel8():
    signature = {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: '*fp32', 4: 'i32'}
    constants = {
        "XBLOCK": 2048
    }
    triton_helper_compile(func=fused_kernel8, signature=signature, constants=constants)

compiled_kernel_list["fused_kernel8"] = test_compile_for_fused_kernel8